# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import hashlib
import hmac
import json
import logging
import time

from tornado import gen
from tornado.websocket import websocket_connect

from coin.exchange.bitmex.ws.table import BitmexTable


class BitmexPrivateSubscriber(object):
  def __init__(self, api_key, api_secret, logger=None):
    self._logger = logger or logging.getLogger(__name__)
    self._api_key = api_key
    self._api_secret = api_secret

    self._ioloop = None
    self._ws_conn = None
    self._stop = False

    self._table = None
    self._table_error_count = 0
    self._restarting = False
    self.callbacks = []  # action, table_name, key, data
    self.on_ready_callback = lambda _: None

  def _init_table(self):
    if self._table is not None:
      self._table.callbacks = []

    self._table = BitmexTable(logger=self._logger)
    self._table.callbacks += self.callbacks
    self._table.error_callback = self._on_table_error
    self._table_error_count = 0

  def _on_table_error(self):
    self._table_error_count += 1
    if self._table_error_count >= 60:
      self._close()
      self._table_error_count = 0

  def _close(self):
    if self._ws_conn is not None:
      try:
        if self._table is not None:
          self._table.callbacks = []
        self._ws_conn.close()
      except Exception:
        self._logger.exception('Error while closing. Ignore error.')
      self._ws_conn = None
      self._table = None

  @gen.coroutine
  def _clean_up_table(self):
    yield gen.sleep(10)
    while not self._stop:
      yield gen.sleep(120)
      if self._table is not None:
        self._table.delete_old_entry()

  @gen.coroutine
  def _ping(self):
    yield gen.sleep(10)
    while not self._stop:
      yield gen.sleep(30)
      if self._ws_conn is not None:
        self._logger.debug('ping')
        try:
          yield self._ws_conn.write_message('ping')
        except Exception:
          self._logger.exception('Ping error')
          if self._restarting is False:
            self._close()
            self.on_ready_callback(False)
            self._logger.error('Connection closed! Reconnecting in 10 seconds...')
            self._ioloop.add_timeout(datetime.timedelta(seconds=10), self._main)
            self._restarting = True

  @gen.coroutine
  def _connect(self):
    assert self._ws_conn is None
    self._ws_conn = yield websocket_connect('wss://www.bitmex.com/realtime')
    res_text = yield self._ws_conn.read_message()
    self._logger.debug('< %s' % res_text)

  @gen.coroutine
  def _auth(self):
    nonce = int(round(time.time()))
    message = 'GET/realtime' + str(nonce)
    signature = hmac.new(bytes(self._api_secret, 'utf8'),
                         bytes(message, 'utf8'),
                         digestmod=hashlib.sha256).hexdigest()
    req_text = json.dumps({'op': 'authKey', 'args': [self._api_key, nonce, signature]})
    yield self._ws_conn.write_message(req_text)
    res_text = yield self._ws_conn.read_message()
    self._logger.debug('< %s' % res_text)
    res = json.loads(res_text)
    assert res['success']

  @gen.coroutine
  def _subscribe(self):
    req_text = json.dumps({'op': 'subscribe', 'args': ['order', 'margin', 'position', 'execution']})
    yield self._ws_conn.write_message(req_text)

  @gen.coroutine
  def _main(self):
    while True:
      try:
        yield self._connect()
        yield self._auth()
        yield self._subscribe()
        break
      except Exception:
        self._logger.exception('Connection failed. Reconnecting in 10 seconds...')
        self._close()
        yield gen.sleep(10)

    self._init_table()
    self.on_ready_callback(True)
    self._restarting = False  # restarting done
    while not self._stop:
      try:
        msg_text = yield self._ws_conn.read_message()
      except Exception:
        self._logger.exception('Error reading ws message')
        self._close()
        msg_text = None

      if msg_text is None:
        if self._restarting is False:
          self._close()
          self.on_ready_callback(False)
          self._logger.error('Connection closed! Reconnecting in 10 seconds...')
          self._ioloop.add_timeout(datetime.timedelta(seconds=10), self._main)
          self._restarting = True
        break

      if msg_text == 'pong':
        self._logger.debug('pong')
        continue

      try:
        msg = json.loads(msg_text)
      except json.decoder.JSONDecodeError:
        self._logger.exception('Failed to decode json. Skipping\n%s' % msg_text)
        continue

      if isinstance(msg, dict):
        self._table.update(msg)
      else:
        self._logger.error('Unknown msg: %s' % str(msg_text))

  def start(self, ioloop):
    assert self._ioloop is None
    self._ioloop = ioloop
    self._ioloop.add_callback(self._main)
    self._ioloop.add_callback(self._ping)
    self._ioloop.add_callback(self._clean_up_table)

  def stop(self):
    self._stop = True


# simple test
def main(_):
  from tornado.ioloop import IOLoop
  subscriber = BitmexPrivateSubscriber(key['access_key'], key['secret_key'])
  ioloop = IOLoop.current()
  subscriber.start(ioloop)
  ioloop.start()


if __name__ == '__main__':
  from absl import app
  app.run(main)
